import os
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from pathlib import Path
import matplotlib.pyplot as plt
import matplotlib
from matplotlib.font_manager import FontProperties

# 设置中文字体
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']  # 优先使用的中文字体
matplotlib.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
font_path = '/Users/liuyuzhuo/Library/Fonts/SimHei.ttf'  # 请确保此路径正确

# 尝试设置中文字体，如果文件存在
if os.path.exists(font_path):
    font_prop = FontProperties(fname=font_path)
    plt.rcParams['font.family'] = font_prop.get_name()
else:
    print("警告：找不到中文字体文件，可能导致中文显示为乱码")
    font_prop = None

def crop_image_remove_title(img_path, title_height_percent=12):
    """裁剪图片，完全去除顶部的主标题区域，但保留列标签
    
    Args:
        img_path: 图片路径
        title_height_percent: 标题区域占总高度的百分比，增加到12%以确保完全去除
    
    Returns:
        裁剪后的图片
    """
    try:
        # 打开图片
        img = Image.open(img_path)
        
        # 计算裁剪区域 - 增加裁剪的高度，确保完全去除标题
        width, height = img.size
        crop_top = int(height * title_height_percent / 100)
        
        # 裁剪图片，完全去除顶部的主标题
        cropped_img = img.crop((0, crop_top, width, height))
        
        return cropped_img
    except Exception as e:
        print(f"处理图片 {img_path} 时出错: {e}")
        return None

def add_column_headers(merged_img):
    """为合并的图像添加列标签
    
    Args:
        merged_img: 合并后的图像
    
    Returns:
        添加了列标签的图像
    """
    width, height = merged_img.size
    
    # 创建一个新图像，为列标签添加空间
    header_height = 40  # 列标签高度
    new_img = Image.new('RGB', (width, height + header_height), (255, 255, 255))
    
    # 粘贴原图像
    new_img.paste(merged_img, (0, header_height))
    
    # 准备绘图
    draw = ImageDraw.Draw(new_img)
    
    # 尝试加载字体
    try:
        font = ImageFont.truetype(font_path, 24) if os.path.exists(font_path) else None
    except:
        font = None
        print("无法加载中文字体，将使用默认字体")
    
    # 计算三列的中心位置
    col_width = width // 3
    col_centers = [col_width // 2, col_width + col_width // 2, 2 * col_width + col_width // 2]
    
    # 定义标签
    labels = ["灰度图像", "真实彩色", "生成彩色"]
    
    # 绘制标签
    for i, label in enumerate(labels):
        # 计算文字大小，以便居中
        if font:
            text_width = draw.textlength(label, font=font)
            text_pos = (col_centers[i] - text_width // 2, 5)
            draw.text(text_pos, label, font=font, fill=(0, 0, 0))
        else:
            # 如果没有字体，使用默认字体
            draw.text((col_centers[i] - 30, 10), label, fill=(0, 0, 0))
    
    return new_img

def merge_images(input_dir, output_file, file_pattern, title):
    """合并多个图片为一个大图，去除重复的主标题，但保留列标签
    
    Args:
        input_dir: 输入目录
        output_file: 输出文件路径
        file_pattern: 文件名模式，用于筛选图片
        title: 标题
    """
    # 获取匹配的文件
    files = [f for f in os.listdir(input_dir) if f.startswith(file_pattern) and f.endswith('.jpg')]
    if not files:
        print(f"没有找到匹配 {file_pattern} 的图片")
        return False
    
    # 排序文件
    files.sort()
    
    print(f"找到 {len(files)} 个匹配的文件: {files}")
    
    # 处理每个图像，只去除主标题区域
    processed_images = []
    for file in files:
        img_path = os.path.join(input_dir, file)
        cropped_img = crop_image_remove_title(img_path)
        if cropped_img:
            processed_images.append(cropped_img)
    
    if not processed_images:
        print("没有有效的图片可以合并")
        return False
    
    # 获取处理后图片的尺寸
    width = processed_images[0].width
    height = processed_images[0].height
    
    # 创建基本合并图像
    merged_height = height * len(processed_images)
    merged_img = Image.new('RGB', (width, merged_height))
    
    # 粘贴每张图片
    for i, img in enumerate(processed_images):
        merged_img.paste(img, (0, i * height))
    
    # 添加列标签
    final_img = add_column_headers(merged_img)
    
    # 保存为临时图片
    temp_path = output_file.replace('.png', '_temp.png')
    final_img.save(temp_path)
    
    # 使用matplotlib添加标题
    plt.figure(figsize=(width/100, (merged_height+50)/100 + 1))
    img_array = np.array(final_img)
    plt.imshow(img_array)
    plt.axis('off')
    
    # 添加总标题
    if font_prop:
        plt.title(title, fontproperties=font_prop, fontsize=26, pad=20)
    else:
        plt.title(title, fontsize=26, pad=20)
    
    # 保存图片
    plt.tight_layout()
    plt.savefig(output_file, dpi=150, bbox_inches='tight')
    plt.close()
    
    # 删除临时文件
    try:
        os.remove(temp_path)
    except:
        pass
    
    print(f"已创建合成图: {output_file}")
    return True

def main():
    # 设置路径
    input_dir = "output/cn_samples"
    output_dir = "output/cn_samples"
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 合并人脸图片
    face_output = os.path.join(output_dir, "face_combined.png")
    merge_images(input_dir, face_output, "face_", "人脸模型着色效果")
    
    # 合并花卉图片
    flower_output = os.path.join(output_dir, "flower_combined.png")
    merge_images(input_dir, flower_output, "flower_", "花卉模型着色效果")
    
    # 更新HTML文件以显示合并后的图片
    html_path = os.path.join(output_dir, "combined.html")
    with open(html_path, 'w', encoding='utf-8') as f:
        f.write("<!DOCTYPE html>\n")
        f.write("<html>\n<head>\n")
        f.write("<meta charset='utf-8'>\n")
        f.write("<title>模型着色效果比较</title>\n")
        f.write("<style>body{font-family:'Microsoft YaHei',Arial,sans-serif;margin:20px;text-align:center;}h1,h2{color:#333;}img{max-width:100%;margin-bottom:30px;}</style>\n")
        f.write("</head>\n<body>\n")
        f.write("<h1>模型着色效果比较</h1>\n")
        
        # 显示合并后的图片
        f.write("<div>\n")
        f.write(f"<h2>人脸模型着色效果</h2>\n")
        f.write(f"<img src='{os.path.basename(face_output)}'>\n")
        f.write("</div>\n")
        
        f.write("<div>\n")
        f.write(f"<h2>花卉模型着色效果</h2>\n")
        f.write(f"<img src='{os.path.basename(flower_output)}'>\n")
        f.write("</div>\n")
        
        f.write("</body>\n</html>")
    
    print(f"处理完成! 已生成合并后的图片和HTML页面")
    print(f"HTML页面: {html_path}")
    print(f"人脸合成图: {face_output}")
    print(f"花卉合成图: {flower_output}")

if __name__ == "__main__":
    main() 